#include "tcp_connection.hh"

#include <iostream>
#define CHECKPOINT std::cout << "CC_CHECKPOINT" << std::endl
// Dummy implementation of a TCP connection

// For Lab 4, please replace with a real implementation that passes the
// automated checks run by `make check`.

template <typename... Targs>
void DUMMY_CODE(Targs &&.../* unused */) {}

using namespace std;
size_t TCPConnection::remaining_outbound_capacity() const { return _sender.stream_in().remaining_capacity(); }

size_t TCPConnection::bytes_in_flight() const { return _sender.bytes_in_flight(); }

size_t TCPConnection::unassembled_bytes() const { return _receiver.unassembled_bytes(); }

size_t TCPConnection::time_since_last_segment_received() const { return _timer - _last_arrived_time; }

void TCPConnection::turn_off_connection() {
    _active = false;
    _linger_after_streams_finish = false;

    _sender.stream_in().set_error();
    _receiver.stream_out().set_error();
    return;
}

void TCPConnection::send_rst(std::optional<WrappingInt32> seqno) {
    TCPSegment rstseg;
    if (seqno.has_value())
        rstseg.header().seqno = seqno.value();
    rstseg.header().rst = true;

    _segments_out.push(rstseg);
    return;
}
void TCPConnection::drive_sender_payload() {
    _sender.fill_window();
    while (!_sender.segments_out().empty()) {
        TCPSegment tmp_seg = std::move(_sender.segments_out().front());
        TCPHeader &tmp_h = tmp_seg.header();
        _sender.segments_out().pop();

        tmp_h.ack = _receiver.ackno().has_value();
        if (_receiver.ackno().has_value()) {
            tmp_h.ackno = _receiver.ackno().value();
        }
        tmp_h.win = _receiver.window_size();
        _segments_out.push(tmp_seg);
    }
}
void TCPConnection::drive_sender_empty() {
    _sender.send_empty_segment();
    TCPSegment tmp_seg = std::move(_sender.segments_out().front());
    TCPHeader &tmp_h = tmp_seg.header();
    _sender.segments_out().pop();

    tmp_h.ack = _receiver.ackno().has_value();
    if (_receiver.ackno().has_value()) {
        tmp_h.ackno = _receiver.ackno().value();
    }
    tmp_h.win = _receiver.window_size();
    _segments_out.push(tmp_seg);
}
void TCPConnection::drive_sender() {
    _sender.fill_window();
    if (_sender.segments_out().empty())
        _sender.send_empty_segment();

    while (!_sender.segments_out().empty()) {
        TCPSegment tmp_seg = std::move(_sender.segments_out().front());
        TCPHeader &tmp_h = tmp_seg.header();
        _sender.segments_out().pop();

        tmp_h.ack = _receiver.ackno().has_value();
        if (_receiver.ackno().has_value()) {
            tmp_h.ackno = _receiver.ackno().value();
            tmp_h.win = _receiver.window_size();
        }
        _segments_out.push(tmp_seg);
    }
}
void TCPConnection::segment_received(const TCPSegment &seg) {
    _last_arrived_time = _timer;
    // is TCP connection still alive ?
    if (!_active)
        return;
    TCPHeader segheader = seg.header();
    // throw std::logic_error("Not implementation\n");
    // segment processs actions

    // first the segment must be acceptable?
    bool acceptable = true;
    // second is this must to send out?
    bool must_send = seg.length_in_sequence_space() > 0;

    // if RST arrives , then tear off the connection
    if (segheader.rst) {
        turn_off_connection();
        return;
    }
    if (!_receiver.ackno().has_value()) {
        // if receiver is not initialized, care about syn and ack
        // std::cout << "LISTEN state or SYN_SENT state" << segheader.syn << std::endl;
        if (!_sender.syn_sent()) {
            // in LISTEN state , ack is rejected but syn is received
            // ack should rejected
            // std::cout << "LISTEN state with coming segment syn :" << segheader.syn << std::endl;
            acceptable = acceptable && !segheader.ack;
            acceptable = acceptable && segheader.syn;
            must_send = must_send && segheader.syn;
        } else {
            // in SYN_SENT state
            uint64_t abs_ackno = unwrap(segheader.ackno, _sender.isn(), _sender.next_seqno_absolute());

            acceptable = acceptable && (!segheader.ack || abs_ackno <= _sender.next_seqno_absolute());
            acceptable = acceptable && segheader.syn;
            must_send = must_send && segheader.syn;
        }
    } else {
        // check the bond
        uint64_t abs_seqno = unwrap(segheader.seqno, _receiver.isn(), _receiver.stream_out().bytes_written());
        uint64_t abs_revnxt =
            unwrap(_receiver.ackno().value(), _receiver.isn(), _receiver.stream_out().bytes_written());
        uint64_t seg_len = seg.length_in_sequence_space();

        bool seq_in_bond = (abs_seqno < abs_revnxt + _receiver.window_size());
        seq_in_bond = seq_in_bond || ((abs_seqno + seg_len - 1) < abs_revnxt + _receiver.window_size());
        seq_in_bond = seq_in_bond || (seg_len == 0 && std::min(abs_seqno, abs_revnxt) == abs_revnxt);
        acceptable = acceptable && seq_in_bond;
        // if( !seq_in_bond ) {
        //     std::cerr   << "seg " << segheader.summary()
        //                 << " is"  << (acceptable ? " not" : "")
        //                 <<  " banned for not in received bond " << std::endl
        //                 << "\tabs_seqno is " << abs_seqno << ",abs_revnxt is " << abs_revnxt
        //                 << ",window size is " << _receiver.window_size()
        //                 << ",size is " << seg.length_in_sequence_space()
        //                 << ",unassemble bytes is " << _receiver.unassembled_bytes() << std::endl;
        // }

        // if the segment " come from past" , now the remote should be announced
        must_send = must_send || (abs_seqno < abs_revnxt && (abs_seqno + seg_len - 1) < abs_revnxt);

        // deal with ack
        uint64_t abs_sndnxtseq = unwrap(segheader.ackno, _sender.isn(), _sender.stream_in().bytes_read());
        bool ack_in_bond_or_not_ack = !segheader.ack || (abs_sndnxtseq <= _sender.next_seqno_absolute());
        acceptable = acceptable && ack_in_bond_or_not_ack;
        // if(!ack_in_bond_or_not_ack) {
        //     std::cerr   << "seg" << segheader.summary()
        //                 << " is" << (acceptable ? " not" : "")
        //                 << " banned for not in ack bond " << std::endl
        //                 << "\tabs ackno is " << abs_sndnxtseq
        //                 << ",next abs seq of sender is " << _sender.next_seqno_absolute() << std::endl;
        // }
    }

    // if acceptable
    if (acceptable)
        _receiver.segment_received(seg);
    if (acceptable && segheader.ack) {
        _sender.ack_received(segheader.ackno, segheader.win);
    }
    // std::cerr   << "segment with " << segheader.summary() << " is"
    //             << (acceptable ? "" : " not") << " acceptable and is"
    //             << (must_send ? ""  : " not") << " must sent" << std::endl;

    // sender process action
    // if acceptable
    //      if must send
    //          fill window and send ack
    //          if no segment have sent out
    //              send empty segment
    // else send a empty ack to show it is unacceptable
    if (acceptable) {
        if (must_send) {
            // take place then must send one
            drive_sender();
        } else {
            drive_sender_payload();
        }
    } else {
        if (must_send)  // not acceptable , but with payload
            drive_sender_empty();
    }

    // judge whether go into lingering
    // FIN_REVD
    if (_receiver.stream_out().input_ended() && TCPState::state_summary(_sender) == TCPSenderStateSummary::SYN_ACKED) {
        _linger_after_streams_finish = false;
        // std::cout << "passived close captured " << std::endl;
    }
    // CLOSING
    if (TCPState::state_summary(_receiver) == TCPReceiverStateSummary::FIN_RECV &&
        TCPState::state_summary(_sender) == TCPSenderStateSummary::FIN_ACKED && !_lingering) {
        _lingering = true;
        _linger_start = _timer;
    }
    // CLOSED
    if (TCPState::state_summary(_receiver) == TCPReceiverStateSummary::FIN_RECV &&
        TCPState::state_summary(_sender) == TCPSenderStateSummary::FIN_ACKED && !_linger_after_streams_finish) {
        _active = false;
    }
}

bool TCPConnection::active() const { return _active; }

size_t TCPConnection::write(const string &data) {
    size_t ret = 0;
    if (!_active)
        return ret;
    if (!_sender.stream_in().input_ended() && !_sender.stream_in().error()) {
        ret = _sender.stream_in().write(data);
    }
    // try to tx data immediately
    drive_sender_payload();
    return ret;
}

//! \param[in] ms_since_last_tick number of milliseconds since the last call to this method
void TCPConnection::tick(const size_t ms_since_last_tick) {
    _timer += ms_since_last_tick;
    _sender.tick(ms_since_last_tick);
    if (_sender.consecutive_retransmissions() > _cfg.MAX_RETX_ATTEMPTS) {
        turn_off_connection();
        send_rst(_receiver.ackno().value());
        return;
    }
    if (!_sender.segments_out().empty()) {
        // retx happen
        TCPSegment retxseg = _sender.segments_out().front();

        retxseg.header().ack = _receiver.ackno().has_value();
        if (_receiver.ackno().has_value()) {
            retxseg.header().ackno = _receiver.ackno().value();
            retxseg.header().win = _receiver.window_size();
        }
        _sender.segments_out().pop();
        // std::cerr << "retx seg :" << retxseg.header().summary() << std::endl;
        _segments_out.push(retxseg);
    }
    if (_lingering && (_timer - _linger_start) >= (10 * _cfg.rt_timeout)) {
        _active = false;
        _lingering = false;
        _linger_after_streams_finish = false;
    }
    return;
}

void TCPConnection::end_input_stream() {
    _sender.stream_in().end_input();
    drive_sender_payload();
}

void TCPConnection::connect() {
    if (_sender.syn_sent())
        throw std::logic_error("The connection is already open or listining to other peer");
    _sender.fill_window();  // get SYN segment
    _segments_out.push(std::move(_sender.segments_out().front()));
    _sender.segments_out().pop();

    return;
}

TCPConnection::~TCPConnection() {
    try {
        if (active()) {
            turn_off_connection();
            send_rst(_sender.next_seqno());
        }
    } catch (const exception &e) {
        std::cerr << "Exception destructing TCP FSM: " << e.what() << std::endl;
    }
}
